use std::collections::{HashMap, HashSet};

use crate::storage::compaction::table_compaction::RemappedRecordLocation;
use crate::storage::index::persisted_bucket_hash_map::GlobalIndex;
use crate::storage::mooncake_table::SnapshotTask;
use crate::storage::mooncake_table_config::MooncakeTableConfig;
use crate::storage::storage_utils::{MooncakeDataFileRef, RecordLocation};

use more_asserts as ma;

/// Struct which stores unpersisted content for mooncake table.
#[derive(Clone, Debug, Default)]
pub(crate) struct UnpersistedRecords {
    /// Mooncake table config.
    mooncake_table_config: MooncakeTableConfig,

    /// ===========================================
    /// Records generated by new writes
    /// ===========================================
    ///
    /// Unpersisted data files, new data files are appended to the end.
    new_data_files: Vec<MooncakeDataFileRef>,
    /// Unpersisted file indices, new indices are appended to the end.
    new_file_indices: Vec<GlobalIndex>,

    /// ===========================================
    /// Records generated by index merge
    /// ===========================================
    ///
    /// Unpersisted old merged file indices, which should not appear in the later iceberg snapshots.
    merged_file_indices_to_remove: Vec<GlobalIndex>,
    /// Unpersisted new merged indices, which should be added to the later iceberg snapshots.
    ///
    /// TODO(hjiang): Consider using hash set for faster lookup.
    merged_file_indices_to_add: Vec<GlobalIndex>,

    /// ===========================================
    /// Records generated by data compaction
    /// ===========================================
    ///
    /// Unpersisted old compacted data files, which should not appear in the later iceberg snapshots.
    compacted_data_files_to_remove: Vec<MooncakeDataFileRef>,
    /// Unpersisted new compacted data files, which should be added in the later iceberg snapshots.
    compacted_data_files_to_add: Vec<MooncakeDataFileRef>,
    /// Unpersisted old compacted file indices, which should not appear in the later iceberg snapshots.
    compacted_file_indices_to_remove: Vec<GlobalIndex>,
    /// Unpersisted new compacted file indices, which should be added in the later iceberg snapshots.
    compacted_file_indices_to_add: Vec<GlobalIndex>,
    /// Un-applied remapped record.
    ///
    /// For a committed deletion record, it could appear in two places: committed deletion log, or iceberg deletion vector puffin blob.
    /// For later, iceberg table manager should handle it, with the knowledge of remapping information.
    compacted_data_file_remap: HashMap<RecordLocation, RemappedRecordLocation>,
}

impl UnpersistedRecords {
    pub(crate) fn new(mooncake_table_config: MooncakeTableConfig) -> Self {
        Self {
            mooncake_table_config,
            ..Default::default()
        }
    }

    /// ==================================
    /// Getters
    /// ==================================
    ///
    pub(crate) fn get_unpersisted_data_files(&self) -> Vec<MooncakeDataFileRef> {
        self.new_data_files.clone()
    }
    pub(crate) fn get_unpersisted_file_indices(&self) -> Vec<GlobalIndex> {
        self.new_file_indices.clone()
    }
    /// Index merge.
    pub(crate) fn get_merged_file_indices_to_add(&self) -> Vec<GlobalIndex> {
        self.merged_file_indices_to_add.clone()
    }
    pub(crate) fn get_merged_file_indices_to_remove(&self) -> Vec<GlobalIndex> {
        self.merged_file_indices_to_remove.clone()
    }
    /// Data compaction.
    pub(crate) fn get_compacted_data_files_to_add(&self) -> Vec<MooncakeDataFileRef> {
        self.compacted_data_files_to_add.clone()
    }
    pub(crate) fn get_compacted_data_files_to_remove(&self) -> Vec<MooncakeDataFileRef> {
        self.compacted_data_files_to_remove.clone()
    }
    pub(crate) fn get_compacted_file_indices_to_add(&self) -> Vec<GlobalIndex> {
        self.compacted_file_indices_to_add.clone()
    }
    pub(crate) fn get_compacted_file_indices_to_remove(&self) -> Vec<GlobalIndex> {
        self.compacted_file_indices_to_remove.clone()
    }
    pub(crate) fn get_compacted_data_file_remap(
        &self,
    ) -> HashMap<RecordLocation, RemappedRecordLocation> {
        self.compacted_data_file_remap.clone()
    }

    /// Get unpersisted data files as hash set for lookup.
    pub(crate) fn get_unpersisted_data_files_set(&self) -> HashSet<MooncakeDataFileRef> {
        let expected_len = self.new_data_files.len() + self.compacted_data_files_to_add.len();
        let mut unpersisted_data_files = HashSet::new();
        unpersisted_data_files.extend(self.new_data_files.iter().cloned());
        unpersisted_data_files.extend(self.compacted_data_files_to_add.iter().cloned());
        assert_eq!(expected_len, unpersisted_data_files.len());
        unpersisted_data_files
    }
    /// Get unpersisted file indices as hash set for lookup.
    #[allow(clippy::mutable_key_type)]
    pub(crate) fn get_unpersisted_file_indices_set(&self) -> HashSet<GlobalIndex> {
        let expected_len = self.new_file_indices.len()
            + self.merged_file_indices_to_add.len()
            + self.compacted_data_files_to_add.len();
        let mut unpersisted_file_indices = HashSet::new();
        unpersisted_file_indices.extend(self.new_file_indices.iter().cloned());
        unpersisted_file_indices.extend(self.merged_file_indices_to_add.iter().cloned());
        unpersisted_file_indices.extend(self.compacted_file_indices_to_add.iter().cloned());
        assert_eq!(expected_len, unpersisted_file_indices.len());
        unpersisted_file_indices
    }

    /// ==================================
    /// Buffer utils
    /// ==================================
    ///
    fn buffer_unpersisted_iceberg_new_data_files(&mut self, task: &SnapshotTask) {
        let new_data_files = task.get_new_data_files();
        self.new_data_files.extend(new_data_files);
    }
    fn buffer_unpersisted_iceberg_new_file_indices(&mut self, task: &SnapshotTask) {
        let new_file_indices = task.get_new_file_indices();
        self.new_file_indices.extend(new_file_indices);
    }
    fn buffer_unpersisted_iceberg_merged_file_indices(&mut self, task: &SnapshotTask) {
        let index_merge_result = &task.index_merge_result;
        self.merged_file_indices_to_add
            .extend(index_merge_result.new_file_indices.to_owned());
        self.merged_file_indices_to_remove
            .extend(index_merge_result.old_file_indices.to_owned());
    }

    fn buffer_unpersisted_iceberg_compaction_data(&mut self, task: &SnapshotTask) {
        let data_compaction_res = &task.data_compaction_result;
        if data_compaction_res.is_empty() {
            return;
        }

        let data_compaction_res = data_compaction_res.clone();
        let l = data_compaction_res.new_data_files.len();

        let mut new_compacted_data_files = Vec::with_capacity(l);
        for (new_data_file, _) in data_compaction_res.new_data_files.into_iter() {
            new_compacted_data_files.push(new_data_file.clone());
        }

        self.compacted_data_files_to_add
            .extend(new_compacted_data_files);
        self.compacted_data_files_to_remove
            .extend(data_compaction_res.old_data_files);
        self.compacted_file_indices_to_add
            .extend(data_compaction_res.new_file_indices);
        self.compacted_file_indices_to_remove
            .extend(data_compaction_res.old_file_indices);

        assert!(self.compacted_data_file_remap.is_empty());
        self.compacted_data_file_remap = data_compaction_res.remapped_data_files;
    }

    /// Buffer unpersisted records.
    pub(crate) fn buffer_unpersisted_records(&mut self, task: &SnapshotTask) {
        self.buffer_unpersisted_iceberg_new_data_files(task);
        self.buffer_unpersisted_iceberg_new_file_indices(task);
        self.buffer_unpersisted_iceberg_merged_file_indices(task);
        self.buffer_unpersisted_iceberg_compaction_data(task);
    }

    /// ==================================
    /// Prune utils
    /// ==================================
    ///
    /// Update unpersisted data files from successful iceberg snapshot operation.
    fn prune_persisted_data_files(&mut self, task: &SnapshotTask) {
        let persisted_new_data_files = &task.persisted_records.import_result.new_data_files;
        ma::assert_ge!(self.new_data_files.len(), persisted_new_data_files.len());
        self.new_data_files.drain(0..persisted_new_data_files.len());
    }

    /// Update unpersisted file indices from successful iceberg snapshot operation.
    fn prune_persisted_file_indices(&mut self, task: &SnapshotTask) {
        let persisted_new_file_indices = &task.persisted_records.import_result.new_file_indices;
        ma::assert_ge!(
            self.new_file_indices.len(),
            persisted_new_file_indices.len()
        );
        self.new_file_indices
            .drain(0..persisted_new_file_indices.len());
    }

    fn prune_persisted_merged_indices(&mut self, task: &SnapshotTask) {
        let old_merged_file_indices = &task
            .persisted_records
            .index_merge_result
            .old_file_indices_removed;
        ma::assert_ge!(
            self.merged_file_indices_to_remove.len(),
            old_merged_file_indices.len()
        );
        self.merged_file_indices_to_remove
            .drain(0..old_merged_file_indices.len());

        let new_merged_file_indices = &task
            .persisted_records
            .index_merge_result
            .new_file_indices_imported;
        ma::assert_ge!(
            self.merged_file_indices_to_add.len(),
            new_merged_file_indices.len()
        );
        self.merged_file_indices_to_add
            .drain(0..new_merged_file_indices.len());
    }

    fn prune_persisted_compacted_data(&mut self, task: &SnapshotTask) {
        let persisted_compaction_res = &task.persisted_records.data_compaction_result;

        ma::assert_ge!(
            self.compacted_data_files_to_add.len(),
            persisted_compaction_res.new_data_files_imported.len()
        );
        self.compacted_data_files_to_add
            .drain(0..persisted_compaction_res.new_data_files_imported.len());

        ma::assert_ge!(
            self.compacted_data_files_to_remove.len(),
            persisted_compaction_res.old_data_files_removed.len()
        );
        self.compacted_data_files_to_remove
            .drain(0..persisted_compaction_res.old_data_files_removed.len());

        ma::assert_ge!(
            self.compacted_file_indices_to_add.len(),
            persisted_compaction_res.new_file_indices_imported.len()
        );
        self.compacted_file_indices_to_add
            .drain(0..persisted_compaction_res.new_file_indices_imported.len());

        ma::assert_ge!(
            self.compacted_file_indices_to_remove.len(),
            persisted_compaction_res.old_file_indices_removed.len()
        );
        self.compacted_file_indices_to_remove
            .drain(0..persisted_compaction_res.old_file_indices_removed.len());

        if !persisted_compaction_res.data_file_records_remap.is_empty() {
            assert_eq!(
                persisted_compaction_res.data_file_records_remap.len(),
                self.compacted_data_file_remap.len()
            );
            self.compacted_data_file_remap.clear();
        }
    }

    /// Prune persisted records.
    pub(crate) fn prune_persisted_records(&mut self, task: &SnapshotTask) {
        self.prune_persisted_data_files(task);
        self.prune_persisted_file_indices(task);
        self.prune_persisted_merged_indices(task);
        self.prune_persisted_compacted_data(task);
    }

    /// ==================================
    /// Persistence utils
    /// ==================================
    ///
    /// Util function to decide whether to create iceberg snapshot by data compaction results.
    fn if_persist_by_data_compaction(&self, force_create: bool) -> bool {
        let new_compacted_data_files_threshold = if !force_create {
            self.mooncake_table_config
                .iceberg_snapshot_new_compacted_data_file_count()
        } else {
            1
        };
        let old_compacted_data_files_threshold = if !force_create {
            self.mooncake_table_config
                .iceberg_snapshot_old_compacted_data_file_count()
        } else {
            1
        };

        self.compacted_data_files_to_add.len() >= new_compacted_data_files_threshold
            || self.compacted_data_files_to_remove.len() >= old_compacted_data_files_threshold
    }

    fn if_persist_by_index_merge(&self, force_create: bool) -> bool {
        let old_merged_file_indices_threshold = if !force_create {
            self.mooncake_table_config
                .iceberg_snapshot_old_merged_file_indices_count()
        } else {
            1
        };

        self.merged_file_indices_to_add.len() >= old_merged_file_indices_threshold
    }

    /// Util function to decide whether to create iceberg snapshot by new data files.
    fn if_persist_by_data_files(&self, force_create: bool) -> bool {
        let data_file_snapshot_threshold = if !force_create {
            self.mooncake_table_config
                .iceberg_snapshot_new_data_file_count()
        } else {
            1
        };
        self.new_data_files.len() >= data_file_snapshot_threshold
    }

    /// Util function to decide whether to flush by new data files or maintenance task.
    pub(crate) fn if_persist_by_new_files_or_maintenance(&self, force_create: bool) -> bool {
        if self.if_persist_by_data_files(force_create) {
            return true;
        }
        if self.if_persist_by_data_compaction(force_create) {
            return true;
        }
        if self.if_persist_by_index_merge(force_create) {
            return true;
        }
        false
    }

    /// ==================================
    /// Validate utils
    /// ==================================
    ///
    /// Validate invariants for persistence buffer, panic if validation fails.
    #[cfg(any(test, debug_assertions))]
    pub(super) fn validate_invariants(&self) {
        // Skip validation for append-only tables since they don't have file indices
        if self.mooncake_table_config.append_only {
            return;
        }

        // Validate new data files and file indices.
        let new_data_files_empty = self.new_data_files.is_empty();
        let new_file_indices_empty = self.new_file_indices.is_empty();
        assert_eq!(new_data_files_empty, new_file_indices_empty);

        // Validate index merge buffer.
        let new_merged_file_indices_empty = self.merged_file_indices_to_add.is_empty();
        let old_merged_file_indices_empty = self.merged_file_indices_to_remove.is_empty();
        assert_eq!(new_merged_file_indices_empty, old_merged_file_indices_empty);

        // Validate data compaction buffer.
        let old_compacted_data_files_empty = self.compacted_data_files_to_remove.is_empty();
        let new_compacted_data_files_empty = self.compacted_data_files_to_add.is_empty();
        let old_compacted_file_indices_empty = self.compacted_data_files_to_remove.is_empty();
        let new_compacted_file_indices_empty = self.compacted_file_indices_to_add.is_empty();
        if old_compacted_data_files_empty {
            assert!(new_compacted_data_files_empty);
            assert!(old_compacted_file_indices_empty);
            assert!(new_compacted_file_indices_empty);
        } else {
            assert!(!old_compacted_file_indices_empty);
        }
    }
}
